import copy
import os
import random
import threading
import cv2
from EasyGui import EasyGui
import tkinter as tk
from tkinter import filedialog
import functools
from ImageReader import ImageReader as IR
from CsvWriter import CsvWriter as CW
import tkinter.messagebox

'数据集自动标注工具2.0'
__author__ = 'littleHuang'

'''GUI界面, 程序由此启动'''


class GUI(EasyGui):
    def __init__(self):
        # GUI的各种绘制
        super().__init__()
        self.setWindowName('数据集自动标注软件2.0')
        self.setLabel('数据集类型', 10, 0, 2, 5)
        self.dataType = tk.StringVar()
        self.setRadiobutton('图像', 'image', self.dataType, 0, 8, 4)
        self.setRadiobutton('视频', 'video', self.dataType, 0, 12, 4)
        self.setLabel('采样间隔', 8, 0, 16, 4)
        self.samplingEntry = self.setEntry(2, 0, 20, 1)
        self.samplingEntry.insert('end', 1)
        self.setLabel('数据集路径', 10, 1, 2, 5)
        self.dataPath = self.setEntry(26, 1, 7, 13)
        self.setButton('选择', 4, self.getDataPath, 1, 20, 2)
        self.labelWay = tk.StringVar()
        self.setLabel('标注方式', 8, 2, 3, 4)
        self.setRadiobutton('手动', 'hand', self.labelWay, 2, 8, 4)
        self.setRadiobutton('自动', 'auto', self.labelWay, 2, 12, 4)
        self.setLabel('训练集：测试集：验证集', 22, 3, 0, 11)
        self.train = self.setEntry(2, 3, 12)
        self.setLabel(':', 1, 3, 13)
        self.test = self.setEntry(2, 3, 14)
        self.setLabel(':', 1, 3, 15)
        self.validation = self.setEntry(2, 3, 16)
        self.setLabel('训练集图像路径', 14, 4, 0, 7)
        self.trainPathEntry = self.setEntry(26, 4, 7, 13)
        self.setButton('选择', 4, functools.partial(self.getPath, entry=self.trainPathEntry), 4, 20, 2)
        self.setLabel('测试集图像路径', 14, 5, 0, 7)
        self.testPathEntry = self.setEntry(26, 5, 7, 13)
        self.setButton('选择', 4, functools.partial(self.getPath, entry=self.testPathEntry), 5, 20, 2)
        self.setLabel('验证集图像路径', 14, 6, 0, 7)
        self.validationPathEntry = self.setEntry(26, 6, 7, 13)
        self.setButton('选择', 4, functools.partial(self.getPath, entry=self.validationPathEntry), 6, 20, 2)
        self.setLabel('训练集csv路径', 14, 7, 0, 7)
        self.trainCsvPathEntry = self.setEntry(26, 7, 7, 13)
        self.setButton('选择', 4, functools.partial(self.getFilePath, entry=self.trainCsvPathEntry), 7, 20, 2)
        self.setLabel('测试集csv路径', 14, 8, 0, 7)
        self.testCsvPathEntry = self.setEntry(26, 8, 7, 13)
        self.setButton('选择', 4, functools.partial(self.getFilePath, entry=self.testCsvPathEntry), 8, 20, 2)
        self.setLabel('验证集csv路径', 14, 9, 0, 7)
        self.validationCsvPathEntry = self.setEntry(26, 9, 7, 13)
        self.setButton('选择', 4, functools.partial(self.getFilePath, entry=self.validationCsvPathEntry), 9, 20, 2)
        self.writeWay = tk.StringVar()
        self.setLabel('csv写入方式', 12, 10, 1, 6)
        self.setRadiobutton('覆盖', 'w', self.writeWay, 10, 8, 4)
        self.setRadiobutton('追加', 'a', self.writeWay, 10, 12, 4)
        self.locType = tk.StringVar()
        self.setLabel('坐标表示方式', 12, 11, 1, 6)
        self.setRadiobutton('像素', 'pixel', self.locType, 11, 8, 4)
        self.setRadiobutton('比例', 'ratio', self.locType, 11, 12, 4)
        self.setLabel('表头设置(输入下列属性的命名, 逗号隔开)', 44, 12, 0, 22)
        self.setLabel('图名,图宽,图高,类型,左上X,左上Y,右下X,右下Y', 44, 13, 0, 22)
        self.titleEntry = self.setEntry(44, 14, 0, 22)
        self.titleEntry.insert(0, 'filename,width,height,class,xmin,ymin,xmax,ymax')
        # 是否正在标注, 防止重复打开多个标注线程
        self.labeling = False
        # 全自动标注工具
        self.autoTrack = False
        # 添加菜单栏
        menubar = tk.Menu(self.fatherWindow)
        statusMenu = tk.Menu()
        menubar.add_cascade(label='状态设置', menu=statusMenu)
        statusMenu.add_command(label='开始', command=self.srartThread)
        statusMenu.add_command(label='退出', command=self.exit)
        statusMenu.add_command(label='全自动模式', command=self.setAutoTrack)
        self.trackType = 'KCF'
        trackTypeChoose = tk.Menu()
        menubar.add_cascade(label='追踪算法', menu=trackTypeChoose)
        for i in ['BOOSTING', 'MIL', 'KCF', 'TLD', 'MEDIANFLOW', 'MOSSE', 'CSRT']:
            trackTypeChoose.add_command(label=i, command=functools.partial(self.setTrackerType, trackerType=i))
        menubar.add_command(label='说明', command=self.instruction)
        self.fatherWindow.config(menu=menubar)
        self.setLabel('消息框', 44, 15, 0, 22)
        self.messageText = self.setText(44, 10, 16, 0, 22)
        self.csvWriter = None
        self.ir = None  # 图像处理类
        fourcc = cv2.VideoWriter_fourcc(*'XVID')
        self.colorWriter = cv2.VideoWriter('final_out.avi',fourcc, 60, (640,480))


        # 设置全自动标注

    def setAutoTrack(self):
        if self.autoTrack:
            self.autoTrack = False
            self.messageText.insert('end', '结束全自动标注\n')
        else:
            self.autoTrack = True
            self.messageText.insert('end', '开始全自动标注\n')

    # 退出命令
    def exit(self):
        self.colorWriter.release()
        try:
            # 关闭csv
            self.csvWriter.close()
        except:
            pass
        # 关闭GUI窗口
        self.fatherWindow.quit()
        self.fatherWindow.destroy()
        # 关闭opencv窗口
        cv2.destroyAllWindows()
        # self.labeling = False
        os._exit(1)  # 退出整个进程, 使用sys.exit(） 只是关闭线程, 会导致线程残留

    # 设置追踪算法
    def setTrackerType(self, trackerType):
        self.trackType = trackerType
        if self.ir:
            self.ir.trackerType = self.trackType
        self.messageText.insert('end', '设置追踪算法为%s\n' % trackerType)

    # 显示说明窗口
    def instruction(self):
        # msg = open('./readme.txt').read()
        # tkinter.messagebox.showinfo(title='说明', message=msg)
        threading.Thread(target=lambda: os.system('gedit readme.txt')).start()

    # 判断路径是否存在
    def isExists(self, path, text):
        if not os.path.exists(path):
            self.messageText.insert('end', text + '路径不存在\n')
            return False
        return True

    # 数据标注主线程
    def start(self):
        dataPath = self.dataPath.get()
        trainImagePath = self.trainPathEntry.get()
        testImagePath = self.testPathEntry.get()
        validationImagePath = self.validationPathEntry.get()
        trainCsvPath = self.trainCsvPathEntry.get()
        testCsvPath = self.testCsvPathEntry.get()
        validationCsvPath = self.validationCsvPathEntry.get()
        message = ['数据集', '训练集图像保存', '测试集图像保存', '验证集图像保存', '训练集csv文件保存', '测试集csv文件保存', '验证集csv文件保存']
        # 检查所有路径是否都存在
        for i, path in enumerate(
                [dataPath, trainImagePath, testImagePath, validationImagePath, trainCsvPath, testCsvPath,
                 validationCsvPath]):
            if not self.isExists(path, message[i]):
                self.labeling = False
                return
        # 判断各个输入框是否都正确填写
        labelWay = self.labelWay.get()
        if labelWay not in ['hand', 'auto']:
            self.messageText.insert('end', '请先选择标注方式\n')
            self.labeling = False
            return
        try:
            trainNum = int(self.train.get())
        except:
            self.messageText.insert('end', '请输入训练集比例\n')
            self.labeling = False
            return
        try:
            testNum = int(self.test.get())
        except:
            self.messageText.insert('end', '请输入测试集比例\n')
            self.labeling = False
            return
        try:
            validationNum = int(self.validation.get())
        except:
            self.messageText.insert('end', '请输入验证集比例\n')
            self.labeling = False
            return
        try:
            sample = int(self.samplingEntry.get())
        except:
            self.messageText.insert('end', '请输入采样比例\n')
            self.labeling = False
            return
        writeWay = self.writeWay.get()
        if writeWay not in ('w', 'a'):
            self.messageText.insert('end', '请选择写入方式\n')
            self.labeling = False
            return
        locType = self.locType.get()
        if locType not in ('pixel', 'ratio'):
            self.messageText.insert('end', '请选择坐标表示方式\n')
            self.labeling = False
            return
        if self.dataType.get() == 'image':
            self.ir = IR(imagePath=dataPath, trackerType=self.trackType)
        elif self.dataType.get() == 'video':
            self.ir = IR(videoPath=dataPath, trackerType=self.trackType)
        else:
            self.messageText.insert('end', '请先选择数据类型')
            self.labeling = False
            return
        # 初始化相机
        if not self.ir.status:
            self.messageText.insert('end', '相机初始化失败\n')
            self.labeling = False
            return
        # 生成可迭代的图像文件读取对象
        imgReader = self.ir.getImage()

        # 初始化csv writer
        try:
            csvWriter = CW(trainCsvPath, testCsvPath, validationCsvPath, writeWay)
            self.csvWriter = csvWriter
        except:
            self.messageText.insert('end', 'csv文件路径有误\n')
            self.labeling = False
            return
        # 用字典存储便于批量处理
        writers = [csvWriter.trw, csvWriter.tew, csvWriter.vw]
        imgPaths = [trainImagePath, testImagePath, validationImagePath]
        # 覆盖方式写入的话需要重现写入表头
        if writeWay == 'w':
            for w in writers:
                w.writerow(self.titleEntry.get().split(','))
        # 采用随机方式按比例划分数据集
        choseList = []
        for i in range(trainNum):
            choseList.append(0)
        for i in range(testNum):
            choseList.append(1)
        for i in range(validationNum):
            choseList.append(2)

        # 数据保存函数
        def save(frame, bboxs, className, imageName):
            # 随机选择写入训练集or测试集or验证集
            index = choseList[random.randint(0, len(choseList) - 1)]
            H = frame.shape[0]
            W = frame.shape[1]
            for i, box in enumerate(bboxs):
                # 获取左上和右下两个坐标点的坐标
                xmin = max(0, int(box[0]))
                ymin = max(0, int(box[1]))
                xmax = min(int(box[0] + box[2]), W)
                ymax = min(int(box[1] + box[3]), H)
                # 按比例方式的话除以图像宽和高
                if locType == 'ratio':
                    xmin = xmin / W
                    ymin = ymin / H
                    xmax = xmax / W
                    ymax = ymax / H
                # csv写入
                writers[index].writerow([imageName, W, H, className[i], xmin, ymin, xmax, ymax])
            # 保存图像
            cv2.imwrite(os.path.join(imgPaths[index], imageName), frame)
            print('图像%s保存成功\n' % imageName)

        # 区分不同目标, 随机产生30种颜色用于画框(如果多于30个目标, 此处可能导致犯错)
        colors = [(random.randint(64, 255), random.randint(64, 255), random.randint(64, 255)) for i in range(30)]
        if writeWay == 'w':
            nameNum = 0
        # 追加方式需要防止文件重名
        else:
            savedNumbers = []
            for dir in [testImagePath, trainImagePath, validationImagePath]:
                for i in os.listdir(dir):
                    try:
                        savedNumbers.append(int(i.split('.')[0]))
                    except:
                        pass
            try:
                nameNum = max(savedNumbers) + 1
            except:
                nameNum = 0

        # 获取ROI函数
        def getBboxs(frame):
            self.messageText.insert('end', '请选择标注框\n')
            while True:
                # 复制一张图像, 避免对源图像的修改
                writeFrame = copy.deepcopy(frame)
                if labelWay == 'hand':
                    bboxs, names = self.ir.getBboxs(writeFrame)
                # 自动方式写需要重置追踪器
                else:
                    bboxs, names = self.ir.initTrackerMulti(writeFrame)
                    # print('tracker init successfully!')
                for i, box in enumerate(bboxs):
                    # print(box)
                    # 展示一下用户绘制的ROI, 并允许重新绘制
                    p1 = (int(box[0]), int(box[1]))
                    p2 = (int(box[0] + box[2]), int(box[1] + box[3]))
                    cv2.rectangle(writeFrame, p1, p2, colors[i], 2, 1)
                    cv2.putText(writeFrame, names[i], p1, cv2.FONT_HERSHEY_SIMPLEX, 1, colors[i], 1)
                # print('show the labeled images')
                self.messageText.insert('end', '按r键重新绘制,j键跳过,其他键保存\n')
                cv2.imshow('showWindows', writeFrame)
                key = cv2.waitKey(0)
                # 按r键重新绘制
                if key == 114:
                    continue
                # 按j键跳过
                elif key == 106:
                    break
                # 保存
                else:
                    save(frame, bboxs, names, str(nameNum) + '.jpg')
                    break
            return names

        if labelWay == 'hand':  # 手动标注
            while True:
                try:
                    for i in range(sample - 1):
                        next(imgReader)
                    frame = next(imgReader)
                    getBboxs(frame)
                    nameNum += 1
                except BaseException as e:
                    # 图像全部读完或者其他错误, 跳出循环
                    print(e)
                    break
        elif labelWay == 'auto':
            try:
                frame = next(imgReader)
                # 初始化追踪器
                names = getBboxs(frame)
                while True:
                    for i in range(sample - 1):
                        self.ir.tracker.update(next(imgReader))
                    print('读取下一张图片')
                    frame = next(imgReader)
                    _, bboxs = self.ir.tracker.update(frame)
                    writeFrame = copy.copy(frame)
                    for i, box in enumerate(bboxs):
                        p1 = (int(box[0]), int(box[1]))
                        p2 = (int(box[0] + box[2]), int(box[1] + box[3]))
                        cv2.rectangle(writeFrame, p1, p2, colors[i], 2, 1)
                        cv2.putText(writeFrame, names[i], p1, cv2.FONT_HERSHEY_SIMPLEX, 1, colors[i], 1)
                    cv2.imshow('showWindows', writeFrame)
                    self.colorWriter.write(writeFrame)
                    if not self.autoTrack:
                        key = cv2.waitKey(0)
                        self.messageText.insert('end', '按r键设定追踪器,j键跳过,其他键保存\n')
                    else:
                        cv2.waitKey(10)
                        key = 100
                    # 按r键重新绘制追踪区域
                    if key == 114:
                        names = getBboxs(frame)
                        nameNum += 1
                    # 按j键跳过
                    elif key == 106:
                        continue
                    else:
                        save(frame, bboxs, names, str(nameNum) + '.jpg')
                        nameNum += 1
            except BaseException as e:
                print(e)
        csvWriter.close()
        cv2.destroyAllWindows()
        self.labeling = False

    # 启动标注线程, 因为图像界面必须有一个线程去维护, 耗时比较长或者需要新开界面的方法最好新开一个线程
    def srartThread(self):
        if self.labeling:
            self.messageText.insert('标注程序已经运行\n')
        t = threading.Thread(target=self.start)
        t.start()

    # 交互获取数据集文件位置
    def getDataPath(self):
        self.dataPath.delete(0, 'end')
        if self.dataType.get() == 'image':
            path = filedialog.askdirectory()
        elif self.dataType.get() == 'video':
            path = filedialog.askopenfilename()
        else:
            self.messageText.insert('end', '请先选择数据集类型\n')
            return
        self.dataPath.insert(0, path)

    # 获取文件夹位置
    def getPath(self, entry):
        entry.delete(0, 'end')
        entry.insert(0, filedialog.askdirectory())

    # 获取文件位置
    def getFilePath(self, entry):
        entry.delete(0, 'end')
        entry.insert(0, filedialog.askopenfilename())


if __name__ == '__main__':
    gui = GUI()
    gui.run()
